package com.example.util;

import java.security.SecureRandom;
import java.util.*;

public class RandomUtils {

    private static final Random RANDOM = new SecureRandom();

    public static void main(String[] args) {
        Map<Boolean, Long> rateMap = new HashMap<>();
        rateMap.put(Boolean.TRUE, 24L);
        rateMap.put(Boolean.FALSE, 100L - 24L);

        // 统计 weightedRandom 结果
        int trueCountWeighted = 0;
        int falseCountWeighted = 0;
        double count = 10000.0;
        for (int i = 0; i < count; i++) {
            Boolean mergeSuccess = RandomUtils.weightedRandom(rateMap);
            if (mergeSuccess) {
                trueCountWeighted++;
            } else {
                falseCountWeighted++;
            }
        }

        System.out.println("------ weightedRandom ------");
        System.out.println("TRUE count: " + trueCountWeighted);
        System.out.println("FALSE count: " + falseCountWeighted);
        System.out.println("TRUE probability: " + (1.0 * trueCountWeighted / count));
        System.out.println("FALSE probability: " + (1.0 * falseCountWeighted / count));

        System.out.println("--------------------------");

        // 统计 randomByCumulativeWeight 结果
        int trueCountCumulative = 0;
        int falseCountCumulative = 0;
        for (int i = 0; i < count; i++) {
            Boolean mergeSuccess = RandomUtils.randomByCumulativeWeight(rateMap);
            if (mergeSuccess) {
                trueCountCumulative++;
            } else {
                falseCountCumulative++;
            }
        }

        System.out.println("------ randomByCumulativeWeight ------");
        System.out.println("TRUE count: " + trueCountCumulative);
        System.out.println("FALSE count: " + falseCountCumulative);
        System.out.println("TRUE probability: " + (1.0 * trueCountCumulative / count));
        System.out.println("FALSE probability: " + (1.0 * falseCountCumulative / count));
    }


    public static <T> T randomByCumulativeWeight(Map<T, Long> weightMap) {
        if (weightMap == null || weightMap.isEmpty()) {
            throw new IllegalArgumentException("权重映射不能为空");
        }

        long totalWeight = weightMap.values().stream()
                .filter(Objects::nonNull)
                .mapToLong(Long::longValue)
                .sum();

        if (totalWeight <= 0) {
            throw new IllegalArgumentException("总权重必须大于0");
        }

        long randomValue = RANDOM.nextInt((int) totalWeight);
        long cumulativeWeight = 0;

        for (Map.Entry<T, Long> entry : weightMap.entrySet()) {
            Long weight = entry.getValue();
            if (weight == null || weight <= 0) continue;

            cumulativeWeight += weight;
            if (randomValue < cumulativeWeight) {
                return entry.getKey();
            }
        }

        // 正常情况下不会到这里，作为兜底返回最后一个元素
        List<T> keys = new ArrayList<>(weightMap.keySet());
        return keys.get(keys.size() - 1);
    }

    /**
     * 在指定范围内随机生成整数 [min, max]
     */
    public static int randomInt(int min, int max) {
        return RANDOM.nextInt((max - min) + 1) + min;
    }


    /**
     * 使用 Alias Method 进行高效随机选择
     */
    public static <T> T weightedRandom(Map<T, Long> weightMap) {
        AliasMethod<T> aliasMethod = new AliasMethod<>(weightMap);
        return aliasMethod.random();
    }

    /**
     * 从 `List` 中随机选择一个元素
     */
    public static <T> T randomFromList(List<T> list) {
        if (list == null || list.isEmpty()) {
            throw new IllegalArgumentException("List cannot be null or empty");
        }
        return list.get(RANDOM.nextInt(list.size()));
    }

    /**
     * Alias Method 随机选择器
     */
    private static class AliasMethod<T> {
        private final List<T> elements;
        private final double[] probabilities;
        private final int[] alias;

        public AliasMethod(Map<T, Long> weightMap) {
            if (weightMap == null || weightMap.isEmpty()) {
                throw new IllegalArgumentException("Weight map cannot be null or empty");
            }

            int size = weightMap.size();
            elements = new ArrayList<>(weightMap.keySet());
            probabilities = new double[size];
            alias = new int[size];

            long totalWeight = weightMap.values().stream().mapToLong(Long::longValue).sum();
            double[] normalizedProbabilities = new double[size];

            int index = 0;
            for (Long weight : weightMap.values()) {
                normalizedProbabilities[index++] = (double) weight / totalWeight * size;
            }

            Deque<Integer> small = new ArrayDeque<>();
            Deque<Integer> large = new ArrayDeque<>();

            for (int i = 0; i < size; i++) {
                if (normalizedProbabilities[i] < 1.0) {
                    small.add(i);
                } else {
                    large.add(i);
                }
            }

            while (!small.isEmpty() && !large.isEmpty()) {
                Integer smallIndex = small.poll();
                Integer largeIndex = large.poll();
                if (smallIndex == null || largeIndex == null) break;

                probabilities[smallIndex] = normalizedProbabilities[smallIndex];
                alias[smallIndex] = largeIndex;

                normalizedProbabilities[largeIndex] -= (1.0 - probabilities[smallIndex]);

                if (normalizedProbabilities[largeIndex] < 1.0) {
                    small.add(largeIndex);
                } else {
                    large.add(largeIndex);
                }
            }

            while (!small.isEmpty()) {
                Integer smallIndex = small.poll();
                if (smallIndex != null) probabilities[smallIndex] = 1.0;
            }
            while (!large.isEmpty()) {
                Integer largeIndex = large.poll();
                if (largeIndex != null) probabilities[largeIndex] = 1.0;
            }
        }

        /**
         * 采样随机元素，时间复杂度 O(1)
         */
        public T random() {
            int column = RANDOM.nextInt(elements.size());
            return RANDOM.nextDouble() < probabilities[column] ? elements.get(column) : elements.get(alias[column]);
        }
    }


}
